#include <Eigen/Dense>
#include <stdint.h>
#include <chrono>
#include "autograd/variable.h"
#include "autograd/variable_test.h"
#include "log.h"
#include "eigen/normal.h"
#include "autograd/graph.h"

namespace  ldl
{
namespace autograd
{
void VariableTest::step29() {
    auto func = [](Variable<double>& x)->Variable<double> {
        auto y = x * x * x * x - double(2)*x.square();
        return y;
    };

    auto gx2 = [](double& x)->double {
        return double(12)*(x*x) - 4;
    };

    auto x = Variable(2.0);
    int64_t iters = 10;

    for(int64_t i = 0;i<iters;i++) {
        LogInfo() << "i: " << i << ", x: " << x;
        auto y = func(x);
        x.cleargrad();
        y.backward();

        LogInfo() << "gx2(x.data()): " << gx2(x.data());
        x.data() -= (x.grad().data()) / gx2(x.data());
    }
}

void VariableTest::step33() {
    auto func = [](Variable<double>& x)->Variable<double> {
        auto y = x * x * x * x - double(2)*x.square();
        return y;
    };

    auto x = Variable(2.0);
    int64_t iters = 10;

    for(int64_t i = 0;i<iters;i++) {
        LogInfo() << "i: " << i << ", x: " << x;
        auto y = func(x);
        x.cleargrad();
        y.backward();
        auto gxdata = x.grad().data();
        auto gx = x.grad();

        x.cleargrad();
        // Graph::plot_dot_graph(x.grad(), true, "x.grad.png");
        gx.backward();
        auto gx2 = x.grad();
        LogInfo() << "gxdata: " << gxdata << ", gx2: " << x.grad();
        x.data() -= (gxdata) / (x.grad().data());
    }

    // {
    //     auto x = Variable<double>(4);
    //     auto y = x * x * x;
    //     y.backward();
    //     Graph::plot_dot_graph(y, true, "y.png");
    //     LogInfo() << "<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<";
    //     LogInfo() << "x.grad(): " << x.grad();
    //     auto gx = x.grad();
    //     Graph::plot_dot_graph(gx, true, "gx.png");
    //     x.cleargrad();
    //     Graph::plot_dot_graph(gx, true, "gx_clearx.png");
    //     gx.backward();
    //     Graph::plot_dot_graph(gx, true, "gx_backward.png");
    //     LogInfo() << "x.grad(): " << x.grad();
    //     LogInfo() << "<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<";
    //     Graph::plot_dot_graph(x.grad(), true, "gx2.png");
    //     // Graph::plot_dot_graph(x.grad());
    // }

    // {
    //     auto x = Variable(2.0);
    //     auto a = x.square();
    //     auto y = a.square() + a.square();
    //     y.backward();

    //     LogInfo() << "y: " << y;
    //     LogInfo() << "x.grad(): " << x.grad();
    // }
}

void VariableTest::eigen_test() {
    {
        auto a = Variable(Eigen::MatrixXd{{1.0, 2.0, 3.0},{4.0, 5.0, 6.0},{7.0, 8.0, 9.0}});
        auto b = Variable(Eigen::MatrixXd{{1.0, 2.0, 3.0},{4.0, 5.0, 6.0},{7.0, 8.0, 9.0}});

        auto c = a + b;
        c.backward();
        LogInfo() << "c: " << c;

        LogInfo() << "a.grad(): " << a.grad();
        LogInfo() << "b.grad(): " << b.grad();
    }

    // 性能测试
    int64_t iters = 1000;
    {
        auto x = Variable(ldl_eigen::Normal::matrix(100,100,0,1.0));
        auto y = x + x;
        auto start = std::chrono::high_resolution_clock::now();

        for(int64_t i = 0;i<iters;i++) {
            y = x + x * x;
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        LogInfo() << "Time taken: " << duration.count() << " ms";
    }

    {
        auto x = ldl_eigen::Normal::matrix(100,100,0,1.0);
        Eigen::MatrixXf y = x + x * x;
        auto start = std::chrono::high_resolution_clock::now();

        for(int64_t i = 0;i<1000;i++) {
            y = x + x * x;
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        LogInfo() << "Time taken: " << duration.count() << " ms";
    }

    {
        auto x = ldl_eigen::Normal::matrix(100,100,0,1.0);
        Eigen::MatrixXf y = x + x * x;
        auto start = std::chrono::high_resolution_clock::now();

        for(int64_t i = 0;i<1000;i++) {
            y = x + x * x;
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        LogInfo() << "Time taken: " << duration.count() << " ms";
    }

    {
        auto x = Variable(ldl_eigen::Normal::matrix(100,100,0,1.0));
        auto y = x + x;
        auto start = std::chrono::high_resolution_clock::now();

        for(int64_t i = 0;i<iters;i++) {
            y = x + x * x;
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        LogInfo() << "Time taken: " << duration.count() << " ms";
    }

    {
        auto x = ldl_eigen::Normal::matrix(100,100,0,1.0);
        Eigen::MatrixXf y = x + x * x;
        auto start = std::chrono::high_resolution_clock::now();

        for(int64_t i = 0;i<1000;i++) {
            y = x + x * x;
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        LogInfo() << "Time taken: " << duration.count() << " ms";
    }

    {
        auto x = Variable(ldl_eigen::Normal::matrix(100,100,0,1.0));
        auto y = x + x;
        auto start = std::chrono::high_resolution_clock::now();

        for(int64_t i = 0;i<iters;i++) {
            y = x + x * x;
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        LogInfo() << "Time taken: " << duration.count() << " ms";
    }

    {
        auto x = ldl_eigen::Normal::matrix(100,100,0,1.0);
        Eigen::MatrixXf y = x + x * x;
        auto start = std::chrono::high_resolution_clock::now();

        for(int64_t i = 0;i<1000;i++) {
            y = x + x * x;
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        LogInfo() << "Time taken: " << duration.count() << " ms";
    }

    {
        auto x = Variable(ldl_eigen::Normal::matrix(100,100,0,1.0));
        auto y = x + x;
        auto start = std::chrono::high_resolution_clock::now();

        for(int64_t i = 0;i<iters;i++) {
            y = x + x * x;
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        LogInfo() << "Time taken: " << duration.count() << " ms";
    }
}
void VariableTest::test()
{
    // 加法
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = Variable(3.5f);
        LogInfo() << "b: " << b;

        auto c = a + b;
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
        LogInfo() << "b.grad(): " << b.grad();
    }

    // 乘法
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = Variable(3.5f);
        LogInfo() << "b: " << b;

        auto c = (a + a) * (b + b);
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
        LogInfo() << "b.grad(): " << b.grad();
    }

    // 减法
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = Variable(3.5f);
        LogInfo() << "b: " << b;

        auto c = (a + a) - (b + b);
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
        LogInfo() << "b.grad(): " << b.grad();
    }

    // 除法
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = Variable(3.5f);
        LogInfo() << "b: " << b;

        auto c = (a + a) / (b + b);
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
        LogInfo() << "b.grad(): " << b.grad();
    }

    // Variable + 数
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = 3.5f;
        LogInfo() << "b: " << b;

        auto c = a + b;
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
    }

    // Variable * 数
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = 3.5f;
        LogInfo() << "b: " << b;

        auto c = (a + a) * (b + b);
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
    }

    // Variable - 数
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = 3.5f;;
        LogInfo() << "b: " << b;

        auto c = (a + a) - (b + b);
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
    }

    // Variable / 数
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = 3.5f;
        LogInfo() << "b: " << b;

        auto c = (a + a) / (b + b);
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
    }

    // 数 + Variable
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = 3.5f;
        LogInfo() << "b: " << b;

        auto c = b + a;
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
    }

    // 数 * Variable
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = 3.5f;
        LogInfo() << "b: " << b;

        auto c = (b + b) * (a + a);
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
    }

    // 数 - Variable
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = 3.5f;;
        LogInfo() << "b: " << b;

        auto c = (b + b) - (a + a);
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
    }

    // 数 / Variable
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = 3.5f;
        LogInfo() << "b: " << b;

        auto c = (b + b) / (a + a);
        LogInfo() << "c: " << c;

        c.backward();

        LogInfo() << "a.grad(): " << a.grad();
    }

    // square
    {
        auto a = Variable(2.5f);
        LogInfo() << "a: " << a;

        auto b = a.square();
        LogInfo() << "b: " << b;

        b.backward();
        LogInfo() << "a.grad(): " << a.grad();
    }

    {
        auto x = Variable(2.0);
        auto a = x.square();
        auto y = a.square() + a.square();
        y.backward();

        LogInfo() << "y: " << y;
        LogInfo() << "x.grad(): " << x.grad();
    }

    // Goldstein-Price函数
    {
        auto goldstein = [](Variable<float> x, Variable<float> y)->Variable<float> {
            auto z = (1.0f + (x + y + 1).square() * (19.0f - 14.0f*x + 3.0f*x.square() - 14.0f*y + 6.0f*x*y+3.0f*y.square())) * \
                (30.0f + (2.0f*x - 3.0f*y).square() * (18.0f - 32.0f*x + 12.0f*x.square() + 48.0f*y - 36.0f*x*y + 27.0f*y.square()));
            return z;
        };
        auto x = Variable<float>(1.0);
        auto y = Variable<float>(1.0);
        auto z = goldstein(x, y);
        z.backward();
        LogInfo() << "z: " << z;
        LogInfo() << "x.grad(): " << x.grad() << ", y.grad(): " << y.grad();
    }

    // step28
    {
        auto rosenbrock = [](Variable<double> x0, Variable<double> x1)->Variable<double> {
            auto y = double(100) * (x1 - x0.square()).square().square() + (x0 - 1.0f).square();
            return y;
        };

        auto x0 = Variable(0.0);
        auto x1 = Variable(2.0);
        double lr = 0.001;
        int64_t iters = 1;
        auto start = std::chrono::high_resolution_clock::now();
        for(int64_t i = 0;i<iters;i++) {
            // LogInfo() << "x0: " << x0 << ", x1: " << x1;

            auto y = rosenbrock(x0,x1);
            x0.cleargrad();
            x1.cleargrad();
            y.backward();

            *(x0.m_data_ptr()->data()) -= lr * (x0.grad().data());
            *(x1.m_data_ptr()->data()) -= lr * (x1.grad().data());
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        LogInfo() << "Time taken: " << duration.count() << " ms";
        LogInfo() << "x0: " << x0 << ", x1: " << x1;
    }
}
}
}